﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Runtime.InteropServices;

public class MFTScanner
{
    private const uint GENERIC_READ = 0x80000000;
    private const int FILE_SHARE_READ = 0x1;
    private const int FILE_SHARE_WRITE = 0x2;
    private const int OPEN_EXISTING = 3;
    private const int FSCTL_ENUM_USN_DATA = 0x900b3;
    private static readonly IntPtr INVALID_HANDLE_VALUE = new IntPtr(-1);
    private IntPtr m_Buffer;
    private int m_BufferSize;

    private IntPtr m_hCJ;

    //// MFT_ENUM_DATA
    [DllImport("kernel32.dll", ExactSpelling = true, SetLastError = true, CharSet = CharSet.Auto)]
    private static extern bool DeviceIoControl(IntPtr hDevice, int dwIoControlCode, ref MFT_ENUM_DATA lpInBuffer,
                                               int nInBufferSize, IntPtr lpOutBuffer, int nOutBufferSize,
                                               ref int lpBytesReturned, IntPtr lpOverlapped);

    [DllImport("kernel32.dll", SetLastError = true, CharSet = CharSet.Auto)]
    private static extern IntPtr CreateFile(string lpFileName, uint dwDesiredAccess, int dwShareMode,
                                            IntPtr lpSecurityAttributes, int dwCreationDisposition,
                                            int dwFlagsAndAttributes, IntPtr hTemplateFile);

    [DllImport("kernel32.dll", ExactSpelling = true, SetLastError = true, CharSet = CharSet.Auto)]
    private static extern Int32 CloseHandle(IntPtr lpObject);

    private void Cleanup()
    {
        if (m_hCJ != IntPtr.Zero)
        {
            // Close the volume handle.
            CloseHandle(m_hCJ);
            m_hCJ = INVALID_HANDLE_VALUE;
        }

        if (m_Buffer != IntPtr.Zero)
        {
            // Free the allocated memory
            Marshal.FreeHGlobal(m_Buffer);
            m_Buffer = IntPtr.Zero;
        }
    }


    public IEnumerable<String> EnumerateFiles(string szDriveLetter)
    {
        try
        {
            var usnRecord = default(USN_RECORD);
            var mft = default(MFT_ENUM_DATA);
            var dwRetBytes = 0;
            int cb = 0;
            var dicFrnLookup = new Dictionary<long, FSNode>();

            // This shouldn't be called more than once.
            if (m_Buffer.ToInt32() != 0)
            {
                throw new Exception("invalid buffer");
            }

            // Assign buffer size
            m_BufferSize = 65536;
            //64KB

            // Allocate a buffer to use for reading records.
            m_Buffer = Marshal.AllocHGlobal(m_BufferSize);

            // correct path
            szDriveLetter = szDriveLetter.TrimEnd('\\');

            // Open the volume handle 
            m_hCJ = CreateFile("\\\\.\\" + szDriveLetter, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, IntPtr.Zero,
                                OPEN_EXISTING, 0, IntPtr.Zero);

            // Check if the volume handle is valid.
            if (m_hCJ == INVALID_HANDLE_VALUE)
            {
                throw new Exception("Couldn't open handle to the volume.");
            }

            mft.StartFileReferenceNumber = 0;
            mft.LowUsn = 0;
            mft.HighUsn = long.MaxValue;

            do
            {
                if (DeviceIoControl(m_hCJ, FSCTL_ENUM_USN_DATA, ref mft, Marshal.SizeOf(mft), m_Buffer, m_BufferSize,
                                    ref dwRetBytes, IntPtr.Zero))
                {
                    cb = dwRetBytes;
                    // Pointer to the first record
                    var pUsnRecord = new IntPtr(m_Buffer.ToInt32() + 8);

                    while ((dwRetBytes > 8))
                    {
                        // Copy pointer to USN_RECORD structure.
                        usnRecord = (USN_RECORD) Marshal.PtrToStructure(pUsnRecord, usnRecord.GetType());

                        // The filename within the USN_RECORD.
                        string FileName =
                            Marshal.PtrToStringUni(new IntPtr(pUsnRecord.ToInt32() + usnRecord.FileNameOffset),
                                                   usnRecord.FileNameLength/2);

                        var bIsFile = !usnRecord.FileAttributes.HasFlag(FileAttributes.Directory);
                        dicFrnLookup.Add(usnRecord.FileReferenceNumber,
                                         new FSNode(usnRecord.FileReferenceNumber, usnRecord.ParentFileReferenceNumber,
                                                    FileName, bIsFile));

                        // Pointer to the next record in the buffer.
                        pUsnRecord = new IntPtr(pUsnRecord.ToInt32() + usnRecord.RecordLength);

                        dwRetBytes -= usnRecord.RecordLength;
                    }

                    // The first 8 bytes is always the start of the next USN.
                    mft.StartFileReferenceNumber = Marshal.ReadInt64(m_Buffer, 0);
                }
                else
                {
                    break; // TODO: might not be correct. Was : Exit Do
                }
            } while (!(cb <= 8));

            // Resolve all paths for Files
            foreach (FSNode oFSNode in dicFrnLookup.Values.Where(o => o.IsFile))
            {
                string sFullPath = oFSNode.FileName;
                FSNode oParentFSNode = oFSNode;

                while (dicFrnLookup.TryGetValue(oParentFSNode.ParentFRN, out oParentFSNode))
                {
                    sFullPath = string.Concat(oParentFSNode.FileName, "\\", sFullPath);
                }
                sFullPath = string.Concat(szDriveLetter, "\\", sFullPath);

                yield return sFullPath;
            }
        }
        finally
        {
            //// cleanup
            Cleanup();
        }
    }

    #region Nested type: FSNode

    private class FSNode
    {
        public readonly string FileName;

        public readonly bool IsFile;
        public readonly long ParentFRN;
        public long FRN;

        public FSNode(long lFRN, long lParentFSN, string sFileName, bool bIsFile)
        {
            FRN = lFRN;
            ParentFRN = lParentFSN;
            FileName = sFileName;
            IsFile = bIsFile;
        }
    }

    #endregion


    #region Nested type: MFT_ENUM_DATA

    [StructLayout(LayoutKind.Sequential)]
    private struct MFT_ENUM_DATA
    {
        public long StartFileReferenceNumber;
        public long LowUsn;
        public long HighUsn;
    }

    #endregion

    #region Nested type: USN_RECORD

    [StructLayout(LayoutKind.Sequential)]
    private struct USN_RECORD
    {
        public readonly int RecordLength;
        public readonly short MajorVersion;
        public readonly short MinorVersion;
        public readonly long FileReferenceNumber;
        public readonly long ParentFileReferenceNumber;
        public readonly long Usn;
        public readonly long TimeStamp;
        public readonly int Reason;
        public readonly int SourceInfo;
        public readonly int SecurityId;
        public readonly FileAttributes FileAttributes;
        public readonly short FileNameLength;
        public readonly short FileNameOffset;
    }

    #endregion
}